import random
import numpy as np
import networkx as nx
import matplotlib.pyplot as plt
from Resources.Node import Node

class DAG(object):
    def __init__(self, task_number, edge_multiple, inner_ccr):
        self.task_number = task_number
        self.edge_multiple = edge_multiple
        self.inner_ccr = inner_ccr
        self.task_list = None
        self.Instructions_boundary = [0, 100]

    def generator(self):

        dag = np.zeros((self.task_number, self.task_number), dtype=int)
        task_list = []
        # 生成dag时，首先保证每一个节点都有后继节点，从而构成连通图
        for i in range(0, self.task_number - 1):  # 不可以取值为 task_number - 1
            j = random.randint(i + 1, self.task_number - 1)  # 可以取值为 task_number - 1
            dag[i, j] = 1

        edge_counter = self.task_number - 1  # 记录总共生成了多少条边
        # 保证每个节点都有一个前继节点，从而保证入口和出口的唯一性
        for j in range(1, self.task_number):
            flag = True
            for e in dag[:, j]:
                if e != 0:
                    flag = False
                    break

            if flag:
                i = random.randint(0, j - 1)
                dag[i, j] = 1
                edge_counter = edge_counter + 1

        # 增加一些边, 可以证明, 此时的 edge_counter 一定小于 2*task_number
        edge_left = self.task_number * self.edge_multiple - edge_counter
        for _ in range(0, edge_left):  # 随机分配剩余的边数
            while True:
                i = random.randint(0, self.task_number - 2)
                j = random.randint(i + 1, self.task_number - 1)
                if dag[i, j] == 0:
                    dag[i, j] = 1
                    break
        # 添加计算时间,区间缩放
        for i in range(self.task_number):
            task_id = i
            runtime = np.random.rand() * (self.Instructions_boundary[1] - self.Instructions_boundary[0])\
                      + self.Instructions_boundary[0]
            total_data_transmission = runtime / self.inner_ccr * 10 * pow(2, 10) * 2 * np.random.rand()  # 将时间转化为数据量
            datas = np.random.random(3)
            datas = datas / sum(datas) * total_data_transmission
            total_input_size = datas[0] + datas[1]
            total_output_size = datas[2]
            children = []
            parents = []
            for j in range(self.task_number):
                if dag[j, i] == 1:
                    parents.append(j)
                if dag[i, j] == 1:
                    children.append(j)
            transmission = np.random.random(len(parents))
            transmission = transmission / sum(transmission) * datas[0]

            trans = dict()
            for k in range(len(parents)):
                trans[parents[k]] = transmission[k]
            t = Node(task_id, children, parents, runtime, total_input_size, total_output_size, trans)
            task_list.append(t)

        self.task_list = task_list


    def output(self):
        print("===================================================")
        print("task_number  | ", end='')
        print(self.task_number)
        print("task_number  | ", end='')
        print(self.task_number)
        print("edge_multiple  | ", end='')
        print(self.edge_multiple)
        print("inner_ccr  | ", end='')
        print(self.inner_ccr)
        print("task_list  | ", end='')
        print()
        for node in self.task_list:
            node.OutputNode()
            print("===================================================")
        return


    def draw(self):

        G = nx.DiGraph()
        for i in self.task_list:
            G.add_node(i.task_id)

        for i in self.task_list:
            for j in i.children:
                G.add_edge(i.task_id, j)
        nx.draw(G, cmap=plt.get_cmap('jet'), with_labels=True, font_weight='bold')
        return G